import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import ResNet18_Weights

class BirdClassifier(nn.Module):
    """
    鸟类分类模型，基于预训练的ResNet18，针对200类鸟类优化
    """
    def __init__(self, num_classes):
        super(BirdClassifier, self).__init__()
        # 启用cudnn基准测试以加速训练
        torch.backends.cudnn.benchmark = True
        # 启用确定性，确保结果可复现
        torch.backends.cudnn.deterministic = True
        
        # 使用ResNet18作为基础模型
        self.model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        
        # 微调冻结策略，只冻结早期层
        # 对于鸟类分类，中后层的特征更重要
        for name, param in self.model.named_parameters():
            if "conv1" in name or "bn1" in name:
                param.requires_grad = False
            else:
                param.requires_grad = True
            
        # 获取全连接层的输入特征数
        n_inputs = self.model.fc.in_features
        
        # 针对200类鸟类分类任务，设计更高效的分类头
        # 使用多层全连接结构，逐步提取和转换特征
        self.model.fc = nn.Sequential(
            nn.Linear(n_inputs, 1024),  # 第一个全连接层
            nn.BatchNorm1d(1024),       # 批归一化提高训练稳定性
            nn.SiLU(),                  # 使用SiLU激活函数，效果优于ReLU
            nn.Dropout(0.3),            # 适度的dropout防止过拟合
            nn.Linear(1024, 2048),      # 第二个全连接层，增加网络容量
            nn.BatchNorm1d(2048),       # 批归一化
            nn.SiLU(),                  # SiLU激活函数
            nn.Dropout(0.3),            # dropout
            nn.Linear(2048, 2048),      # 第三个全连接层，进一步增加网络容量
            nn.BatchNorm1d(2048),       # 批归一化
            nn.SiLU(),                  # SiLU激活函数
            nn.Dropout(0.3),            # dropout
            nn.Linear(2048, num_classes) # 输出层
        )
    
    def forward(self, x):
        return self.model(x)

def count_parameters(model):
    """
    计算模型中可训练参数的数量
    """
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def get_model(num_classes):
    """
    创建并返回模型实例，针对200类鸟类分类任务优化
    """
    model = BirdClassifier(num_classes)
    
    # 打印可训练参数数量，用于调试
    print(f"可训练参数数量: {count_parameters(model)}")
    return model